
from light_training.preprocessing.preprocessors.preprocessor_aortaseg2024 import DefaultPreprocessor 
import numpy as np 
import pickle 
import json 

base_dir = "/vip_media/data/miccai24/AortaSeg24/segmamba_code/data/raw_data"

def process_train():
    preprocessor = DefaultPreprocessor(base_dir=base_dir)

    out_spacing = [1, 1, 1]
    output_dir = "./data/fullres/train/"

    with open("./data_analysis_result.txt", "r") as f:
        content = f.read().strip("\n")
        print(content)
    content = eval(content)
    foreground_intensity_properties_per_channel = content["intensity_statistics_per_channel"]
    
    preprocessor.run(output_spacing=out_spacing, 
                     output_dir=output_dir, 
                     all_labels=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
                     num_processes=12,
                     foreground_intensity_properties_per_channel=foreground_intensity_properties_per_channel)


def plan():
    # base_dir = "/vip_media/data/miccai24/AortaSeg24/segmamba_code/data/raw_data"

    preprocessor = DefaultPreprocessor(base_dir=base_dir,  
                                   )
    preprocessor.run_plan()

if __name__ == "__main__":
    # plan()

    process_train()

    # process_test()